package start.aai.common.Ollama;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;

import java.util.*;

@Slf4j
public class OllamaClient {
    private static final WebClient webClient = WebClient.create();
    private static final String OLLAMA_API_URL = "http://localhost:11434/api/chat";
    private static final String DEFAULT_MODEL = "qwen2.5vl:3b";
    public static final int MAX_HISTORY_SIZE = 20;
    public static  List<Map<String, Object>> chatHistory = new ArrayList<>();
    private static final ObjectMapper objectMapper = new ObjectMapper();

    public static Flux<String> callModelStream(List<Map<String, Object>> messages) {
        log.info("===== 调用Ollama模型 =====");

        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", DEFAULT_MODEL);
        requestBody.put("messages", messages);
        requestBody.put("stream", true);

        return webClient.post()
                .uri(OLLAMA_API_URL)
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(requestBody)
                .retrieve()
                .bodyToFlux(String.class)
                // 关键：使用map而非mapNotNull，确保不丢失任何响应（包括错误）
                .mapNotNull(jsonStr -> {
                    log.debug("原始响应: {}", jsonStr.substring(0, Math.min(100, jsonStr.length())));
                    try {
                        JsonNode root = objectMapper.readTree(jsonStr);
                        JsonNode messageNode = root.get("message");
                        boolean isDone = root.get("done").asBoolean(false);

                        // 处理结束标志
                        if (isDone) {
                            return null; // 结束标志不返回内容
                        }

                        // 处理正常消息
                        if (messageNode != null) {
                            String content = messageNode.get("content").asText("").trim();
                            // 过滤空内容
                            if (content.isEmpty()) {
                                return null;
                            }
                            // 清洗data:前缀
                            String cleanContent = content
                                    .replaceAll("(?i)\\bdata:?\\s*", "")
                                    .replaceAll("(?i)(data:?\\s*)+", "")
                                    .trim();
                            log.debug("清洗后内容: {}", cleanContent);
                            return cleanContent;
                        }

                        // 无message节点的异常响应
                        log.warn("Ollama响应无message字段: {}", jsonStr);
                        return "服务响应格式异常，请重试";

                    } catch (Exception e) {
                        // 解析异常时返回具体错误，而非null
                        log.error("解析Ollama响应失败", e);
                        return "服务解析异常: " + e.getMessage().substring(0, Math.min(50, e.getMessage().length()));
                    }
                })
                // 只过滤null（结束标志），保留错误信息
                .filter(Objects::nonNull)
                // 顶层异常捕获
                .onErrorResume(e -> {
                    log.error("Ollama调用整体异常", e);
                    return Flux.just("服务调用失败: " + e.getMessage());
                });
    }

    public static void trimHistory() {
        while (chatHistory.size() > MAX_HISTORY_SIZE) {
            chatHistory.removeFirst();
        }
    }

    public static void clearHistory() {
        chatHistory.clear();
    }
}
